﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Xml;

using System.Diagnostics.Contracts;

namespace SimpleIR
{
	public sealed class IRBody
	{
		private readonly List<IRBasicBlock> m_blocks;
		private readonly Dictionary<long, IRBasicBlock> m_blockIDToEntryMap;

		private IRT1T2Block m_T1T2Rep;

		[ContractInvariantMethod]
		private void ObjectInvariant()
		{
			Contract.Invariant(this.m_blocks != null);
			Contract.Invariant(this.m_blockIDToEntryMap != null);
		}

		public IRBody(List<IRBasicBlock> stmtl)
		{
			Contract.Requires(stmtl != null);

			this.m_blocks = stmtl;

			this.m_blockIDToEntryMap = new Dictionary<long, IRBasicBlock>();
			foreach(IRBasicBlock bb in this.m_blocks)
				this.m_blockIDToEntryMap.Add(bb.BlockID, bb);
		}

		public override String ToString()
		{
			StringBuilder sb = new StringBuilder();
			foreach(IRBasicBlock bb in this.m_blocks)
				sb.Append(bb.ToString() + "\n");
			return sb.ToString();
		}

		public string DGMLDisplayString()
		{
			StringBuilder sbnodes = new StringBuilder();
			StringBuilder sblinks = new StringBuilder();

			System.Xml.XmlWriterSettings settings = new System.Xml.XmlWriterSettings { Indent = true, OmitXmlDeclaration = true, CloseOutput = true, Encoding = Encoding.UTF8, ConformanceLevel = System.Xml.ConformanceLevel.Fragment };
			
			using(XmlWriter xwn = XmlWriter.Create(sbnodes, settings), xwl = XmlWriter.Create(sblinks, settings))
			{
				this.m_T1T2Rep.DGMLDisplayString(xwn, xwl, true);
			}

			System.Xml.XmlWriterSettings settings2 = new System.Xml.XmlWriterSettings { Indent = true, OmitXmlDeclaration = true, CloseOutput = true, Encoding = Encoding.UTF8 };
			StringBuilder sbgraph = new StringBuilder();
			using(XmlWriter xw = XmlWriter.Create(sbgraph, settings2))
			{
				xw.WriteStartDocument();
				xw.WriteStartElement("DirectedGraph", "http://schemas.microsoft.com/vs/2009/dgml");

				xw.WriteStartElement("Nodes");
				xw.WriteRaw(sbnodes.ToString());
				xw.WriteEndElement();

				xw.WriteStartElement("Links");
				xw.WriteRaw(sblinks.ToString());
				xw.WriteEndElement();

				xw.WriteStartElement("Styles");
				xw.WriteRaw("\t<Style TargetType=\"Node\"><Setter Property=\"Style\" Value=\"Plain\" /><Setter Property=\"NodeRadius\" Value=\"2\" /></Style>\n");
				xw.WriteRaw("\t<Style TargetType=\"Node\">\n\t\t<Condition Expression=\"Group = 'Expanded' and HasCategory('Summary')\" />\n\t\t<Setter Property=\"Background\" Value=\"GhostWhite\" />\n\t</Style>\n");
				xw.WriteRaw("\t<Style TargetType=\"Node\">\n\t\t<Condition Expression=\"Group = 'Collapsed'\" />\n\t\t<Setter Property=\"Background\" Value=\"Silver\" />\n\t</Style>\n");
				xw.WriteEndElement();

				xw.WriteEndElement();
				xw.WriteEndDocument();
			}

			return sbgraph.ToString();
		}

        /// <summary>
		/// Get the IRT1T2Block that makes up the body, must have previously called the T1T2 constructor routines.
		/// </summary>
		public IRT1T2Block T1T2Body 
		{ 
			get 
			{
				Contract.Assert(this.m_T1T2Rep != null, "Hmm, this shouldn't be used before calling T1T2Build!!");

				return this.m_T1T2Rep; 
			} 
		}

		/// <summary>
		/// Get the list of all basic blocks (in some order) that make up the method.
		/// DO NOT MODIFY!!!
		/// </summary>
		public List<IRBasicBlock> BasicBlockList
		{ get { return this.m_blocks; } }

		/// <summary>
		/// Return the block with the given id.
		/// </summary>
		[Pure]
		public IRBasicBlock GetBlockWithID(long blckid)
		{ return this.m_blockIDToEntryMap[blckid]; }

		/// <summary>
		/// Return the entry block for this method.
		/// </summary>
		public IRBasicBlock GetMethodEntryBlock()
		{ return this.m_blockIDToEntryMap[0]; }

		//////////////////////////////////////////////////////////////////////////////////////////////////
		//Simplify Methods

		/// <summary>
		/// Add all the static constructor calls to the main method
		/// </summary>
		internal void AddStaticConstructorCalls(List<IRStaticInvoke> scons, IRType voidtype)
		{
			Contract.Requires(scons != null);
			Contract.Requires(voidtype != null);

			List<IRByteCode> calls = new List<IRByteCode>();
			foreach(IRStaticInvoke iid in scons)
			{
				IRCallOp tcall = IRCallOp.MakeStaticMethodCall(voidtype, iid, new IRSourceLocation("StaticInitializer", -1, -1));
				calls.Add(tcall);
			}

			IRBasicBlock entryblock = this.m_blockIDToEntryMap[0];
			entryblock.AddByteCodeListAtStartOfBlock(calls);
		}

		//////////////////////////////////////////////////////////////////////////////////////////////////
		//Build T1-T2 forms and post process

		/// <summary>
		/// Build T1-T2 forms of methods and number the statements for use in the analysis
		/// </summary>
		public void BuildT1T2()
		{
			Dictionary<long, HashSet<long>> loopinfo = this.ComputeLoopHeadersAndBodies();
			HashSet<long> insideloops = new HashSet<long>();
			List<IRT1T2Region> toplvl = new List<IRT1T2Region>();

			this.BuildT1T2ProcBlock(0, toplvl, loopinfo, insideloops, new HashSet<long>());
			toplvl.Reverse();
			this.m_T1T2Rep = new IRT1T2Block(toplvl);
		}

		/// <summary>
		/// The rec walk to build the T1T2 set (int regions) that corresponds to the node block to proc.
		/// Since we are doing this in PostOrder, we add to the back of regions and we have reverse topo sorted the blocks for free.
		/// </summary>
		private void BuildT1T2ProcBlock(long blocktoproc, List<IRT1T2Region> regions, Dictionary<long, HashSet<long>> loopinfo, HashSet<long> insideloops, HashSet<long> visited)
		{
			visited.Add(blocktoproc);

			List<uint> trgts = new List<uint>();
			this.m_blockIDToEntryMap[blocktoproc].GetAllBranchTargets(trgts);
			foreach(uint jt in trgts)
			{
				if(!visited.Contains(jt))
					BuildT1T2ProcBlock(jt, regions, loopinfo, insideloops, visited);
			}

			if(loopinfo.ContainsKey(blocktoproc))
				regions.Add(BuildT1T2Loop(blocktoproc, regions, loopinfo, insideloops));
			else
				regions.Add(this.m_blockIDToEntryMap[blocktoproc]);
		}

		/// <summary>
		/// If we have a block that is a loop head take the needed chunks out of regions build them into a loop and replace the head entry in regions 
		/// with the new loop.
		/// </summary>
		private IRT1T2Loop BuildT1T2Loop(long loophead, List<IRT1T2Region> regions, Dictionary<long, HashSet<long>> loopinfo, HashSet<long> insideloops)
		{
			List<IRT1T2Region> body = new List<IRT1T2Region>();

			HashSet<long> inthisbody = loopinfo[loophead];
			for(int i = 0; i < regions.Count; ++i)
			{
				IRT1T2Region rr = regions[i];
				long rrblockval = (rr is IRBasicBlock ? ((IRBasicBlock)rr).BlockID : ((IRT1T2Loop)rr).Head.BlockID);

				if(insideloops.Contains(rrblockval) || (!inthisbody.Contains(rrblockval))) //already nested in another loop or not in this loop
					continue;

				regions[i] = null;
				body.Add(rr);
			}

			insideloops.UnionWith(inthisbody);
			regions.RemoveAll((IRT1T2Region tr) => tr == null);
			body.Reverse();
			
			return new IRT1T2Loop(this.m_blockIDToEntryMap[loophead], (body.Count > 0 ? new IRT1T2Block(body) : null), new HashSet<long>(loopinfo[loophead]));
		}

		/// <summary>
		/// Compute the set of loop headers in the body and the sets of blocks that make up the loop body (disjoint from the head blocks, so body may be empty)
		/// </summary>
		private Dictionary<long, HashSet<long>> ComputeLoopHeadersAndBodies()
		{
			Dictionary<long, HashSet<long>> predmap = this.ComputePredMap();
			Dictionary<long, HashSet<long>> domins = this.ComputeDominatorSets(predmap);

			Dictionary<long, HashSet<long>> loopheadtobody = new Dictionary<long, HashSet<long>>();

			foreach(IRBasicBlock bbl in this.m_blocks)
			{
				foreach(long pred in predmap[bbl.BlockID])
				{
					if(domins[pred].Contains(bbl.BlockID)) //if the head dominates the tail
					{
						if(!loopheadtobody.ContainsKey(bbl.BlockID))
							loopheadtobody.Add(bbl.BlockID, new HashSet<long>());
						HashSet<long> lbdy = AddAllContainedBlockIDsToLoopBodySet(bbl.BlockID, pred, predmap);
						loopheadtobody[bbl.BlockID].UnionWith(lbdy);
					}
				}
			}

			return loopheadtobody;
		}

		/// <summary>
		/// Add all the bodies between loophead (but not including the head) and the loopback (including this) to the loopbody set.
		/// </summary>
		private static HashSet<long> AddAllContainedBlockIDsToLoopBodySet(long loophead, long loopback, Dictionary<long, HashSet<long>> predmap)
		{
			//taken from ACDI p. 192, m = loopback, n = loophead

			Stack<long> tstack = new Stack<long>();
			HashSet<long> loop = new HashSet<long>();
			loop.Add(loopback);
			loop.Add(loophead); //remove this at the end

			if(loophead != loopback)
				tstack.Push(loopback);

			while(tstack.Count != 0)
			{
				long p = tstack.Pop();

				foreach(long q in predmap[p])
				{
					if(!loop.Contains(q))
					{
						loop.Add(q);
						tstack.Push(q);
					}
				}
			}

			loop.Remove(loophead);
			return loop;
		}

		/// <summary>
		/// Compute the dominator sets for this method body.
		/// </summary>
		/// <returns></returns>
		private Dictionary<long, HashSet<long>> ComputeDominatorSets(Dictionary<long, HashSet<long>> predmap)
		{
			//
			//Impl. taken from Advanced Compiler Design and Implelentation p. 182
			//
			Dictionary<long, HashSet<long>> domin = new Dictionary<long, HashSet<long>>();
			HashSet<long> N = new HashSet<long>(this.m_blockIDToEntryMap.Keys);
			foreach(IRBasicBlock bbl in this.m_blocks)
			{
				HashSet<long> doms = new HashSet<long>();
				if(bbl.BlockID == 0)
					doms.Add(bbl.BlockID);
				else
					doms.UnionWith(N);
				domin.Add(bbl.BlockID, doms);
			}

			bool change = true;
			while(change)
			{
				change = false;
				foreach(IRBasicBlock bb in this.m_blocks)
				{
					if(bb.BlockID == 0)
						continue;

					HashSet<long> T = new HashSet<long>(N);
					foreach(long p in predmap[bb.BlockID])
						T.IntersectWith(domin[p]);

					HashSet<long> D = T;
					D.Add(bb.BlockID);
					if(!domin[bb.BlockID].SetEquals(D))
					{
						domin[bb.BlockID] = D;
						change = true;
					}
				}
			}

			return domin;
		}

		/// <summary>
		/// Compute basic info about what happens in this body.
		/// </summary>
		/// <param name="usedmodsf">The set of statics that may be accessed in some way in this body.</param>
		/// <param name="possibledelgateinst">A map from delegate types to the invoke identities that may be placed in them (null if we cannot precisely resolve the targets).</param>
		internal void InitUsedStaticsDelegateTargetInfo(HashSet<IRStaticField> usedmodsf, Dictionary<IRLambdaType, HashSet<IRInvokeIdentity>> possibledelgateinst)
		{
			Contract.Requires(usedmodsf != null);
			Contract.Requires(possibledelgateinst != null);

			IRInvokeIdentity liid = null;
			IRLambdaType dt = null;
			foreach(IRBasicBlock bb in this.m_blocks)
			{
				foreach(IRByteCode ss in bb.Operations)
				{
					switch(ss.OperationCode)
					{
						case OpCode.Op.AllocCons:
							if(((IRAllocationOp)ss).AllocType is IRLambdaType)
							{
								dt = (IRLambdaType)((IRAllocationOp)ss).AllocType;
								if(!possibledelgateinst.ContainsKey(dt))
									possibledelgateinst.Add(dt, new HashSet<IRInvokeIdentity>());
								
								possibledelgateinst[dt].Add(liid);
								liid = null;
							}
							break;
						case OpCode.Op.LoadFuncPtr:
							Contract.Assert(liid == null, "We lost a step somewhere.");
							liid = ((IRLoadToEvalStackBasic)ss).InvokeIdentity;
							break;
						case OpCode.Op.LoadVFuncPtr:
							throw new NotImplementedException();
						case OpCode.Op.LoadGlobalAddr:
							usedmodsf.Add(((IRPointerIndexOp)ss).StaticField);
							break;
						case OpCode.Op.LoadFromGlobal:
							usedmodsf.Add(((IRLoadToEvalStackBasic)ss).StaticField);
							break;
						case OpCode.Op.StoreToGlobal:
							usedmodsf.Add(((IRStoreFromEvalStackBasic)ss).StaticField);
							break;
						default:
							break;
					}
				}
			}
		}

		/// <summary>
		/// Compute the initial call graph, exclude any call targets we cannot resolve (marking the caller nodes as such).
		/// </summary>
		/// <param name="caller">The InvokeIdentity of this method body.</param>
		/// <param name="cg">The call graph which we want to fill in.</param>
		/// <param name="possibledelgateinst">A map from delegate types to the invoke identities that may be placed in them.</param>
		public void BuildCallGraph(IRInvokeIdentity caller, CallGraph cg, Dictionary<IRLambdaType, HashSet<IRInvokeIdentity>> possibledelgateinst, IRProgram p)
		{
			Contract.Requires(caller != null);
			Contract.Requires(cg != null);
			Contract.Requires(possibledelgateinst != null);
			Contract.Requires(p != null);

			HashSet<IRInvokeIdentity> ptrgts = new HashSet<IRInvokeIdentity>();
			IRCallOp invbc;
			IRAllocationOp invalloc;
			HashSet<IRInvokeIdentity> iivid = null;
			
			foreach(IRBasicBlock bb in this.m_blocks)
			{
				foreach(IRByteCode ss in bb.Operations)
				{
					switch(ss.OperationCode)
					{
						case OpCode.Op.Call:
							invbc = (IRCallOp)ss;
							ptrgts.Add(p.InvokeTable.LookupInvokeIdentity(invbc.SpecificTargetClass, invbc.InvokeSig));
							break;
						case OpCode.Op.CallVirt:
							invbc = (IRCallOp)ss;
							if(invbc.ReciverExpType is IRPtrType && ((IRPtrType)invbc.ReciverExpType).PtrTargetType is IRLambdaType && invbc.InvokeSig.InvokeName.Equals("Invoke"))
							{
								IRLambdaType rctt = (IRLambdaType)((IRPtrType)invbc.ReciverExpType).PtrTargetType;
								if(possibledelgateinst.ContainsKey(rctt))
								{
									iivid = possibledelgateinst[rctt];
									ptrgts.UnionWith(iivid);
								}
								iivid = null;
							}
							else
							{
								Tuple<IRType, IRInvokeSig> nettypeandvsig = new Tuple<IRType, IRInvokeSig>(invbc.BaseImplementationType, invbc.InvokeSig);
								ptrgts.UnionWith(p.GetPossibleInvocationTargetsForVirtCallExp(invbc.ReciverExpType, nettypeandvsig));
							}
							break;
						case OpCode.Op.AllocCons:
							invalloc = (IRAllocationOp)ss;
							if(!invalloc.IsArrayAlloc)
								ptrgts.Add(p.InvokeTable.LookupInvokeIdentity(invalloc.AllocType, invalloc.ConsIdentity.Sig));
							break;
						default:
							break;
					}
				}
			}

			foreach(IRInvokeIdentity trgtcall in ptrgts)
				cg.AddStaticResolvedCallEdgeAsNeeded(caller, trgtcall);
		}

		/// <summary>
		/// Compute the pred. relation on the blocks in the body.
		/// </summary>
		private Dictionary<long, HashSet<long>> ComputePredMap()
		{
			Dictionary<long, HashSet<long>> blockPredMap = new Dictionary<long, HashSet<long>>();
			foreach(IRBasicBlock bbl in this.m_blocks)
				blockPredMap.Add(bbl.BlockID, new HashSet<long>());

			List<uint> jtrgs = new List<uint>();
			foreach(IRBasicBlock bbl in this.m_blocks)
			{
				jtrgs.Clear();
				bbl.GetAllBranchTargets(jtrgs);
				foreach(uint jt in jtrgs)
					blockPredMap[jt].Add(bbl.BlockID);
			}

			return blockPredMap;
		}

		internal void NumberBody(CtrGen cg)
		{
			Contract.Requires(cg != null);

			this.m_T1T2Rep.NumberRegion(cg); 
		}

		internal void DoLiveness(bool isconstructor)
		{
			//get all address taken locals and initialize live after map
			HashSet<IRVarName> addrtakenlocals = new HashSet<IRVarName>();
			Dictionary<long, HashSet<IRVarName>> liveafter = new Dictionary<long, HashSet<IRVarName>>();
			foreach(IRBasicBlock bb in this.m_blocks)
			{
				liveafter.Add(bb.BlockID, new HashSet<IRVarName>());

				foreach(IRByteCode bc in bb.Operations)
				{
					if(bc.OperationCode == OpCode.Op.LoadLocalAddr)
						addrtakenlocals.Add(((IRPointerIndexOp)bc).LocalName);
				}
			}

			if(isconstructor)
				addrtakenlocals.Add(this.GetBlockWithID(0).VarsInScopeAtEntry.First((IRVarName v) => v.NameString.Equals("this")));

			this.m_T1T2Rep.PropLiveVariables(liveafter, addrtakenlocals);
		}

		internal void ComputeSymbolicEvalStack(IRProgram p)
		{
			Dictionary<long, SymbolicEvalStack> evalonenter = new Dictionary<long, SymbolicEvalStack>();
			foreach(IRBasicBlock bb in this.m_blocks)
				evalonenter.Add(bb.BlockID, new SymbolicEvalStack());

			this.m_T1T2Rep.ComputeEvalStackExpressions(evalonenter, p);
		}

		internal void ComputeLoopTraverseInfo(IRProgram p)
		{
			foreach(IRT1T2Region rr in this.m_T1T2Rep.T1T2Blocks)
			{
				bool titerok = true;
				IRVarName titerator = null;
				HashSet<IRCallOp> titergets = new HashSet<IRCallOp>();
				IRCallOp tenumadv = null;
				bool tindexok = true;
				IRVarName tindexvar = null;
				IRMemberField toptfield = null;
				HashSet<IRByteCode> tindexgets = new HashSet<IRByteCode>();
				IRVarName tindexpos = null;
				IRStoreFromEvalStackBasic tindexstore = null;

				rr.DoLoopCopyAnalysis(true, -1, -1, p, ref titerok, ref titerator, titergets, ref tenumadv, ref tindexok, ref tindexvar, ref toptfield, tindexgets, ref tindexpos, ref tindexstore, this.m_blockIDToEntryMap);
			}
		}
	}
}
